// ElGamal implementation by Mike Austin

#include <stdio.h>
#include <stdlib.h>
#include <time.h>
#include <sys/time.h>
#include <unistd.h>
#include <math.h>
#include <string.h>
#include "ElGamal.hpp"
#include "cryptoUtil.h"
#include "wlan.hpp"

#define GET_ITH_BIT( exponent, index )	( ( exponent & ( 1 << index ) ) >> index )
#define X_MOD_POWER_2( x, power ) ( x & ( power - 1 ) )

const static int logging = 0;
static unsigned int thisPublicKey = 0;
static unsigned int thisPrivateKey = 0;

// list of peers, their states and shared key
#define MAX_PEERS 10
static Peer peers[MAX_PEERS];
static unsigned int peerIndex = 0;
bool toggle = false;

struct timeval startTime;
struct timeval endTime;

void log( char * str )
{
	if ( logging )
	{
		printf("%s\n", str);
	}
}

unsigned int calculateGCD(unsigned int a, unsigned int b)
{
	if ( b == 0 )
	{
		return a;
	}
	else
	{
		// integer division is ok here because even if the decimal part gets chopped off,
		// that's what we wanted anyway from the floor() function

		return calculateGCD( b, a - b * /*floor(*/ ( a / b ) /*)*/ );
	}
}

unsigned int computeRandomNumber()
{
	char str[100];
	unsigned int number;
	unsigned int gcd;

	srand((unsigned int)time(NULL));

	struct num
	{
		short part1;
		short part2;
	} * combinedNum = (struct num *) &number;	

	toggle = !toggle;

	do
	{
		// need to calculate random upper 16 bits and random lower 16 bits
		// separately because rand() only goes up to 2^16-1
		if (toggle) {
			combinedNum->part1 = rand() % (prime - 1) + 1;
			combinedNum->part2 = rand() % (prime - 1) + 1;
		}
		else {
			combinedNum->part2 = rand() % (prime - 1) + 1;
			combinedNum->part1 = rand() % (prime - 1) + 1;
		}
		
		sprintf(str, "Calculated random number %u", number);
		log(str);

		gcd = calculateGCD( number, (prime - 1) );
	}
	while ( gcd != 1 );

	sprintf(str, "GCD(%u, %u) = 1", number, (prime - 1));
	log(str);

	return number;
}

unsigned int montgomeryProduct( unsigned int a, unsigned int b )
{
	unsigned int t = a * b;

	unsigned int u = ( t + ( X_MOD_POWER_2( t * montgomery_N, montgomery_R ) * prime ) / montgomery_R );
	if ( u >= prime )
	{
		return u - prime;
	}
	else
	{
		return u;
	}
}

unsigned int montgomeryReduction( unsigned int base, unsigned int exponent, unsigned int modulus )
{
	// Montgomery modular exponentiation

	unsigned int result = 0;
	unsigned int i;

	unsigned int a = ( base * montgomery_R ) % modulus;
	unsigned int x = 5;		// R mod prime = 2^32 mod 2^32 - 5 = 5

	for ( i = 31; i >= 0; i-- )		// i = # bits in exponent minus 1
	{
		x = montgomeryProduct( x, x );
		if ( GET_ITH_BIT( exponent, i ) == 1 )
		{
			x = montgomeryProduct( x, a );
		}
	}
	result = montgomeryProduct( x, 1 );

	return result;
}

unsigned int exponentiationBySquaring( unsigned int base, unsigned int exponent, unsigned int modulus )
{
	unsigned long long result = 1;
	unsigned long long baseL = base;

	while ( exponent > 0 )
	{		
		if ( ( exponent & 1 ) == 1 )
		{
			result = ( result * baseL ) % modulus;
		}
		exponent = exponent >> 1;
		baseL = ( baseL * baseL ) % modulus;
		//printf("result = %llu\n", result );
	}

	if ( result > 0xFFFFFFFF )
	{
		printf("Warning, result is larger than 32 bits\n");
	}

	return result;
}

unsigned int modularExponentiation( unsigned int base, unsigned int exponent, unsigned int modulus )
{
	char str[100];
	sprintf(str, "%u^%u mod %u = ", base, exponent, modulus);
	log(str);

	unsigned int result;

	//result = montgomeryReduction( base, exponent, modulus );
	result = exponentiationBySquaring( base, exponent, modulus );

	sprintf(str, "%u", result);
	log(str);
	
	return result;
}

unsigned int generatePrivateKey()
{
	return computeRandomNumber();
}

unsigned int generatePublicKey( unsigned int a )
{
	return modularExponentiation( generator, a, prime );
}

Ciphertext elGamalEncrypt( unsigned int publicKey, unsigned int data )
{
	unsigned int number;		// ephemeral key
	unsigned int sharedSecret;
	unsigned int c1;
	unsigned int c2;
	unsigned long long multiplier;
	Ciphertext ciphertext = {0};

	printf("elGamalEncrypt plaintext: %u\n", data);

	if ( publicKey > 0 )
	{
		if ( data >= prime )
		{
			printf("Warning: data is greater than prime, taking data mod prime instead\n");
			data %= prime;
		}

		srand((unsigned int)time(NULL));

		struct num
		{
			short part1;
			short part2;
		} * combinedNum = (struct num *) &number;

		// need to calculate random upper 16 bits and random lower 16 bits
		// separately because rand() only goes up to 2^16-1
		toggle = !toggle;
		if (toggle) {
			combinedNum->part1 = rand() % (prime - 1) + 1;
			combinedNum->part2 = rand() % (prime - 1) + 1;
		}
		else {
			combinedNum->part2 = rand() % (prime - 1) + 1;
			combinedNum->part1 = rand() % (prime - 1) + 1;
		}

		c1 = modularExponentiation( generator, number, prime );
		sharedSecret = modularExponentiation( publicKey, number, prime );
		multiplier = (unsigned long long) data * sharedSecret;
		c2 = multiplier % prime;

		ciphertext.c1 = c1;
		ciphertext.c2 = c2;
	}

	printf("elGamalEncrypt ciphertext: (%u,%u)\n", ciphertext.c1, ciphertext.c2);

	return ciphertext;
}

unsigned int elGamalDecrypt( unsigned int privateKey, Ciphertext ciphertext )
{
	unsigned int sharedSecret;
	unsigned int sharedSecretInverse;
	unsigned int data;
	unsigned long long multiplier;

	printf("elGamalDecrypt ciphertext: (%u,%u)\n", ciphertext.c1, ciphertext.c2);

	if ( privateKey > 0 )
	{
		sharedSecret = modularExponentiation( ciphertext.c1, privateKey, prime );
		sharedSecretInverse = modularExponentiation( sharedSecret, prime - 2, prime );
		multiplier = (unsigned long long) ciphertext.c2 * sharedSecretInverse;
		data = multiplier % prime;
	}

	printf("elGamalDecrypt plaintext: %u\n", data);

	return data;
}

int compareMacAddrs( unsigned char * pBuffer1, unsigned char * pBuffer2 )
{
//	int ret = 0;
//	int i;

	/*
	WLANAddr addr1;
	memcpy( addr1.data, pBuffer1, 6 );
	WLANAddr addr2;
	memcpy( addr2.data, pBuffer2, 6 );
	printf("Comparing %s and %s\n", addr1.wlan2asc(), addr2.wlan2asc());

	for ( i = 0; i < 6; i++ )
	{
		if ( pBuffer1[i] != pBuffer2[i] )
		{
			ret = 1;
			break;
		}
	}
	*/

	return memcmp( pBuffer1, pBuffer2, 6 );
}

void sendCryptoPacket( PCryptoPacket pCryptoPacket, WLANHeader * pHdr )
{
	MySendData( pHdr->destAddr.data, pCryptoPacket );
}

void handleElGamal( PCryptoPacket pCryptoPacket, WLANHeader * pHdr )
{
	unsigned int i, number;
	PCryptoPacket pcp = NULL;
	double timeDiff;
	unsigned int messageType = pCryptoPacket->id & 0x0F;
   unsigned char temp;

	//printf("messageType=%u\n", messageType);
	switch ( messageType )
	{
	case ELGAMAL_PUBLIC_KEY_REQUEST:
	{
		printf("Received ELGAMAL_PUBLIC_KEY_REQUEST message from %s\n", pHdr->srcAddr.wlan2asc() );
		gettimeofday( &startTime, NULL );

		// The fact that we're receiving this message means we're not the initiator, so clear the table
		// in case anything went bad previously
		clearElGamalPeerTable();

		// This is an incoming public key request, so send back our public key
		if ( thisPublicKey == 0 || thisPrivateKey == 0 )
		{
			thisPrivateKey = generatePrivateKey();
			thisPublicKey = generatePublicKey( thisPrivateKey );
		}

		// make a peer entry since we're exchanging keys apparently
		if ( peerIndex < MAX_PEERS )
		{
			for ( i = 0; i < WLAN_ADDR_LEN; i++ )
			{
				peers[ peerIndex ].macAddr[i] = pHdr->srcAddr.data[i];
			}
			peers[ peerIndex ].state = ELGAMAL_EXCHANGING_KEY;
			peers[ peerIndex ].publicKey = 0;	// don't have their public key
			peers[ peerIndex ].sharedKey = 0;	// they'll send us a secret key later

			printf("Added new peer in ELGAMAL_EXCHANGING_KEY state: addr %s\n", pHdr->srcAddr.wlan2asc() );
			peerIndex++;
		}

		// send public key response
		pcp = (CryptoPacket *)malloc(sizeof(CryptoPacket));
		if ( pcp != NULL )
		{
			// flip the addresses
			for ( i = 0; i < WLAN_ADDR_LEN; i++ )
			{
            temp = pHdr->destAddr.data[i];
				pHdr->destAddr.data[i] = pHdr->srcAddr.data[i];
				pHdr->srcAddr.data[i] = temp;
			}

			// alg type and message type
			pcp->id = ELGAMAL_PUBLIC_KEY_RESPONSE | ( ELGAMAL << 4 );


			// add the data (this node's public key)
			memcpy( pcp->data, &thisPublicKey, sizeof(unsigned int) );

			// send it
			sendCryptoPacket( pcp, pHdr );
			printf("Sent ELGAMAL_PUBLIC_KEY_RESPONSE with public key %u to %s\n", thisPublicKey, pHdr->destAddr.wlan2asc());
		}
	
		break;
	}
	case ELGAMAL_PUBLIC_KEY_RESPONSE:
	{
		Ciphertext ciphertext = {0};

		printf("Received ELGAMAL_PUBLIC_KEY_RESPONSE message from %s\n", pHdr->srcAddr.wlan2asc() );

		// This is an answer to our public key request, so store the info
		// TODO: Should probably check whether the peer exists already
		if ( peerIndex < MAX_PEERS )
		{
			for ( i = 0; i < WLAN_ADDR_LEN; i++ )
			{
				peers[ peerIndex ].macAddr[i] = pHdr->srcAddr.data[i];
			}
			peers[ peerIndex ].state = ELGAMAL_EXCHANGING_KEY;
			memcpy( &(peers[ peerIndex ].publicKey), pCryptoPacket->data, sizeof(unsigned int) );

			printf("Added new peer in ELGAMAL_EXCHANGING_KEY state: addr %s - public key %u\n", 
				pHdr->srcAddr.wlan2asc(), 
				peers[ peerIndex ].publicKey );

			// now generate a shared secret key and send it
			struct num
			{
				short part1;
				short part2;
			} * combinedNum = (struct num *) &number;	

			// need to calculate random upper 16 bits and random lower 16 bits
			// separately because rand() only goes up to 2^16-1
			toggle = !toggle;
			if (toggle) {
				combinedNum->part1 = rand() % (prime - 1) + 1;
				combinedNum->part2 = rand() % (prime - 1) + 1;
			}
			else {
				combinedNum->part2 = rand() % (prime - 1) + 1;
				combinedNum->part1 = rand() % (prime - 1) + 1;
			}

			// remember the secret key
			peers[ peerIndex ].sharedKey = number;

			// then send it
			pcp = (CryptoPacket *)malloc(sizeof(CryptoPacket));
			if ( pcp != NULL )
			{
				// flip the addresses
				for ( i = 0; i < WLAN_ADDR_LEN; i++ )
				{
					temp = pHdr->destAddr.data[i];
				   pHdr->destAddr.data[i] = pHdr->srcAddr.data[i];
				   pHdr->srcAddr.data[i] = temp;
				}

				// message type
				pcp->id = ELGAMAL_SEND_DATA | ( ELGAMAL << 4 );

				// encrypt the secret key
				ciphertext = elGamalEncrypt( peers[ peerIndex ].publicKey, number );

				// add the data (the encrypted generated secret key)
				memcpy( pcp->data, &ciphertext, sizeof(Ciphertext) );

				// send it
				sendCryptoPacket( pcp, pHdr );
				printf("Sent ELGAMAL_SEND_DATA with shared secret key %u to %s\n", number, pHdr->destAddr.wlan2asc());
			}
			peerIndex++;
		}
		break;
	}
	case ELGAMAL_SEND_DATA:
	{
		printf("Received ELGAMAL_SEND_DATA message from %s\n", pHdr->srcAddr.wlan2asc() );

		unsigned int plaintext = 0;
		Ciphertext ciphertext = {0};
		int found = 0;	

		// We're receiving an encrypted message (probably a shared key)
		if ( thisPrivateKey > 0 )
		{
			memcpy( &ciphertext, pCryptoPacket->data, sizeof(Ciphertext) );
			plaintext = elGamalDecrypt( thisPrivateKey, ciphertext );
			if ( plaintext > 0 )
			{
				printf("Decrypted shared key: %u\n", plaintext);

				// update peer table
				for ( i = 0; i < peerIndex; i++ )
				{
					if ( compareMacAddrs( pHdr->srcAddr.data, peers[i].macAddr ) == 0 )
					{
						peers[i].sharedKey = plaintext;
						peers[i].state = ELGAMAL_READY;		// ready since we have a secret key now
						printf("Changed peer to ELGAMAL_READY state: addr %s - public key %u - shared key %u\n", 
							pHdr->srcAddr.wlan2asc(), 
							peers[ i ].publicKey, plaintext );

						// send ack message
						pcp = (CryptoPacket *)malloc(sizeof(CryptoPacket));
						if ( pcp != NULL )
						{
							// flip the addresses
							for ( i = 0; i < WLAN_ADDR_LEN; i++ )
							{
								temp = pHdr->destAddr.data[i];
				            pHdr->destAddr.data[i] = pHdr->srcAddr.data[i];
				            pHdr->srcAddr.data[i] = temp;
							}

							// message type
							pcp->id = ELGAMAL_DATA_ACK | ( ELGAMAL << 4 );

							// don't need to add any data so send it
							sendCryptoPacket( pcp, pHdr );
							printf("Sent ELGAMAL_DATA_ACK to %s\n", pHdr->destAddr.wlan2asc());
							gettimeofday( &endTime, NULL );
							timeDiff = (double)(endTime.tv_usec - startTime.tv_usec);
							timeDiff /= 1000000;
							printf("Key exchange completed in %f seconds\n", timeDiff );
						}
						found = 1;
						break;
					}
				}
				if ( found == 0 )
				{
					printf("Could not find peer entry for %s\n", pHdr->srcAddr.wlan2asc());
				}
			}
			else
			{
				printf("decryption returned 0, something went wrong\n");
			}
		}
		else
		{
			printf("We don't have a private key for some reason?!\n");
		}
		break;
	}
	case ELGAMAL_DATA_ACK:
	{
		printf("Received ELGAMAL_DATA_ACK message from %s\n", pHdr->srcAddr.wlan2asc() );

		// We're receiving an ack that Bob got the data we sent, update the peer table to ready state
		for ( i = 0; i < peerIndex; i++ )
		{
			if ( compareMacAddrs( pHdr->srcAddr.data, peers[i].macAddr ) == 0 )
			{
				if ( peers[i].sharedKey > 0 )
				{
					peers[i].state = ELGAMAL_READY;
					printf("Changed peer to ELGAMAL_READY state: addr %s - public key %u - shared key %u\n", 
						pHdr->srcAddr.wlan2asc(), 
						peers[ i ].publicKey, peers[ i ].sharedKey );
					gettimeofday( &endTime, NULL );
					timeDiff = (double)(endTime.tv_usec - startTime.tv_usec);
					timeDiff /= 1000000;
					printf("Key exchange completed in %f seconds\n", timeDiff );
				}
			}
		}
		break;
	}
	}
}

void initiateElGamal()
{
	int i;
	WLANHeader hdr;
	PCryptoPacket pcp = (PCryptoPacket) malloc( sizeof(CryptoPacket) );

	gettimeofday( &startTime, NULL );

	// This was initiated by a user menu option, so start everything fresh
	clearElGamalPeerTable();

	for ( i = 0; i < WLAN_ADDR_LEN; i++ )
	{
		hdr.destAddr.data[i] = 0xFF;
	}

	if ( pcp != NULL )
	{
		pcp->id = ELGAMAL_PUBLIC_KEY_REQUEST | ( ELGAMAL << 4 );
		// don't need to send any data
		sendCryptoPacket( pcp, &hdr );
		printf("Sent ELGAMAL_PUBLIC_KEY_REQUEST to ff:ff:ff:ff:ff:ff\n");
	}
}

void printElGamalPeerTable()
{
	int i;
	WLANAddr addr;
	char str[30];

	printf("Peer table\n==================================================================\nAddress\t\t\tPublic key\tState\t\tShared key\n");

	for ( i = 0; i < peerIndex; i++ )
	{
		memcpy( addr.data, peers[i].macAddr, WLAN_ADDR_LEN );
		if ( peers[i].state == ELGAMAL_EXCHANGING_KEY )
		{
			sprintf(str, "ELGAMAL_EXCHANGING_KEY");
		}
		else if ( peers[i].state == ELGAMAL_WAITING )
		{
			sprintf(str, "ELGAMAL_WAITING");
		}
		else if ( peers[i].state == ELGAMAL_READY )
		{
			sprintf(str, "ELGAMAL_READY");
		}
		printf("%s\t%u\t%s\t%u\n", addr.wlan2asc(), peers[i].publicKey, str, peers[i].sharedKey);
	}
}

void clearElGamalPeerTable()
{
	int i, j;
	for ( i = 0; i < MAX_PEERS; i++ )
	{
		for ( j = 0; j < WLAN_ADDR_LEN; j++ )
		{
			peers[i].macAddr[j] = 0;
		}
		peers[i].publicKey = 0;
		peers[i].state = 0;
		peers[i].sharedKey = 0;
	}
	peerIndex = 0;
}

